package com.vislardo.app.notebook.utils;

import com.vislardo.app.notebook.error.FileEncryptorException;

import javax.crypto.*;
import javax.crypto.spec.GCMParameterSpec;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.file.*;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;

public class FileEncryptor {

    // AES-GCM 参数
    private static final String ALGORITHM = "AES/GCM/NoPadding";
    private static final int GCM_IV_LENGTH = 12;  // GCM 推荐 12 字节 IV
    private static final int GCM_TAG_LENGTH = 16; // 128 位认证标签
    private static final int BUFFER_SIZE = 4096;  // 缓冲区大小（字节）

    /**
     * 生成随机密钥（保存此密钥用于后续解密）
     */
    public static SecretKey generateKey() throws NoSuchAlgorithmException {
        KeyGenerator keyGen = KeyGenerator.getInstance("AES");
        keyGen.init(256); // AES-256
        return keyGen.generateKey();
    }

    /**
     * 加密文件（使用 NIO）
     */
    public static void encryptFile(SecretKey key, Path inputPath, Path outputPath)
            throws IOException, FileEncryptorException, NoSuchPaddingException, NoSuchAlgorithmException, InvalidAlgorithmParameterException, InvalidKeyException {
        SecureRandom secureRandom = new SecureRandom();
        byte[] iv = new byte[GCM_IV_LENGTH];
        secureRandom.nextBytes(iv); // 生成随机 IV

        // 初始化加密器
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH * 8, iv);
        cipher.init(Cipher.ENCRYPT_MODE, key, gcmSpec);

        try (FileChannel inChannel = FileChannel.open(inputPath, StandardOpenOption.READ);
             FileChannel outChannel = FileChannel.open(outputPath,
                     StandardOpenOption.CREATE, StandardOpenOption.WRITE)) {

            // 将 IV 写入文件头部
            ByteBuffer ivBuffer = ByteBuffer.wrap(iv);
            outChannel.write(ivBuffer);

            // 分配直接缓冲区提升性能
            ByteBuffer buffer = ByteBuffer.allocateDirect(BUFFER_SIZE);
            while (inChannel.read(buffer) != -1) {
                buffer.flip();
                byte[] data = new byte[buffer.remaining()];
                buffer.get(data);
                byte[] encrypted = cipher.update(data, 0, buffer.limit());
                if (encrypted != null) {
                    outChannel.write(ByteBuffer.wrap(encrypted));
                }
                buffer.clear();
            }

            // 处理最后一个数据块
            byte[] finalEncrypted = cipher.doFinal();
            if (finalEncrypted != null) {
                outChannel.write(ByteBuffer.wrap(finalEncrypted));
            }
        } catch (IllegalBlockSizeException | BadPaddingException e) {
            throw new FileEncryptorException(e.getMessage());
        }
    }


    public static SecretKey encryptFile(Path inputPath, Path outputPath) throws FileEncryptorException,IOException{
        try{
            SecretKey key = generateKey();
            encryptFile(key,inputPath,outputPath);
            return key;
        } catch (InvalidAlgorithmParameterException | NoSuchAlgorithmException| NoSuchPaddingException|InvalidKeyException e) {
            throw new FileEncryptorException(e.getMessage());
        }
    }
    /**
     * 解密文件（使用 NIO）
     */
    public static void decryptFile(SecretKey key, Path inputPath, Path outputPath) throws FileEncryptorException,IOException {
        try (FileChannel inChannel = FileChannel.open(inputPath, StandardOpenOption.READ);
             FileChannel outChannel = FileChannel.open(outputPath,
                     StandardOpenOption.CREATE, StandardOpenOption.WRITE)) {
            // 读取 IV
            ByteBuffer ivBuffer = ByteBuffer.allocate(GCM_IV_LENGTH);
            inChannel.read(ivBuffer);
            ivBuffer.flip();
            byte[] iv = new byte[GCM_IV_LENGTH];
            ivBuffer.get(iv);

            // 初始化解密器
            Cipher cipher = Cipher.getInstance(ALGORITHM);
            GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH * 8, iv);
            cipher.init(Cipher.DECRYPT_MODE, key, gcmSpec);

            // 分配直接缓冲区
            ByteBuffer buffer = ByteBuffer.allocateDirect(BUFFER_SIZE);
            while (inChannel.read(buffer) != -1) {
                buffer.flip();
                byte[] data = new byte[buffer.remaining()];
                buffer.get(data);
                byte[] decrypted = cipher.update(data, 0, buffer.limit());
                if (decrypted != null) {
                    outChannel.write(ByteBuffer.wrap(decrypted));
                }
                buffer.clear();
            }
            // 处理最后一个数据块
            byte[] finalDecrypted = cipher.doFinal();
            if (finalDecrypted != null) {
                outChannel.write(ByteBuffer.wrap(finalDecrypted));
            }
        } catch (InvalidAlgorithmParameterException|NoSuchPaddingException|NoSuchAlgorithmException|InvalidKeyException e) {
            throw new FileEncryptorException(e.getMessage());
        } catch (IllegalBlockSizeException |BadPaddingException e) {
            throw new FileEncryptorException(e.getMessage());
        }
    }
}